package com.example.tree.heap;

import java.util.Comparator;
import java.util.PriorityQueue;
import java.util.Queue;

/**
 * 给你一个 n x n 矩阵 matrix ，其中每行和每列元素均按升序排序，找到矩阵中第 k 小的元素。(和LeetCode23 类似)
 * 请注意，它是 排序后 的第 k 小元素，而不是第 k 个 不同 的元素。
 *
 *  示例 1：
 * 输入：matrix = [[1,5,9],[10,11,13],[12,13,15]], k = 8
 * 输出：13
 * 解释：矩阵中的元素为 [1,5,9,10,11,12,13,13,15]，第 8 小元素是 13
 *
 *  示例 2：
 * 输入：matrix = [[-5]], k = 1
 * 输出：-5
 */
public class Leetcode378_KthSmallest {
    public static void main(String[] args) {
        int[][] matrix = {
                {1,  5,  9},
                {10, 11, 13},
                {12, 13, 15}};

        System.out.println(new Solution().kthSmallest(matrix, 5));

    }

    static class Solution {

        /**
         * 解法二:归并排序的归的过程
         * 一般归并排序是两个数组归并，而本题是 n 个数组归并，所以需要用小根堆维护，以优化找到最小值的时间复杂度。
         *
         * @param matrix
         * @param k
         * @return
         */
        private int kthSmallest2(int[][] matrix, int k) {
            // int[] 是三个元素的数组, num[0]: matrix[i][j], num[1]: i, num[2]: j
//            Queue<int[]> minHeap = new PriorityQueue<>((o1, o2) -> o1[0] - o2[0]);
            Queue<int[]> minHeap = new PriorityQueue<>(Comparator.comparingInt(o -> o[0])); //两种写法等价

            // 将第一列的所有数据加入堆中
            for (int i = 0; i < matrix.length; i++) {
                minHeap.offer(new int[] {matrix[i][0], i, 0});
            }

            // 先合并到 k-1 个元素
            for (int i = 0; i < k - 1; i++) {
                int[] tmp = minHeap.poll();
                if (tmp[2] != matrix.length - 1)  // 没到一行的最后一个元素就将当前元素的后一个(列+1)的元素放入堆中
                    minHeap.offer(new int[]{matrix[tmp[1]][tmp[2] + 1], tmp[1], tmp[2] + 1});
            }

            return minHeap.poll()[0];
        }

        /**
         * 解法一:最大堆
         * @param matrix
         * @param k
         * @return
         */
        private int kthSmallest1(int[][] matrix, int k) {
            Queue<Integer> maxHeap = new PriorityQueue<>(k, (i1, i2) -> i2 - i1);

            for (int i = 0; i < matrix.length; i++) {
                for (int j = 0; j < matrix[0].length; j++) {
                    if (maxHeap.size() < k || maxHeap.peek() > matrix[i][j]) {
                        if (maxHeap.size() == k) maxHeap.poll();

                        maxHeap.offer(matrix[i][j]);
                    }
                }
            }

            return maxHeap.peek();

        }

        public int kthSmallest(int[][] matrix, int k) {
            return kthSmallest2(matrix, k);
        }
    }
}
